#include <stdio.h>

#include "proxy.h"

int glog_level = LOG_ERR; //LOG_INFO; //LOG_ERR;
struct config_s gconfig;
uv_loop_t* gloop = NULL;
int gChilds_num = 0;
char* debug_url = "assets.cnblogs.com:443";

BOOL WINAPI ConsoleCtrlHandler(DWORD dwCtrlType)
{
	if (dwCtrlType == CTRL_C_EVENT) {
		uv_stop(gloop);
		printf("got user stop signal: ctrl+c\n");
		printf("Please wait patiently for a while...\n");
	}
	return TRUE;
}

//CONNECT ntp.msn.cn:443 HTTP/1.1\r\n
int cb_url(llhttp_t* ph, const char* at, size_t length)
{
	if (length >= 255) {
		log_message(LOG_CRIT, "http first line: url len more than 255: %d", length);
		return 0;
	}

	struct llhttps* pllhttp = (struct llhttps*)llhttp_get_priv(ph);
	char* url = pllhttp->url;
	strncpy(url, at, length);
	url[length] = 0;

	pllhttp->found_url = 1;
	log_message(LOG_INFO, "http first line: CONNECT %s", url);
	return 0;
}

int cb_complete(llhttp_t* ph)
{
	struct llhttps* pllhttp = (struct llhttps*)llhttp_get_priv(ph);
	pllhttp->http_header_finish = 1;
	//log_message(LOG_INFO, "reach to http last line");
	return 0;
}

llhttp_settings_t gllhttp_set;

void init_llhttp_set(void)
{
	llhttp_settings_init(&gllhttp_set);
	gllhttp_set.on_url = cb_url;
	gllhttp_set.on_message_complete = cb_complete;
}

//init llhttp
void init_llhttp(struct llhttps* pllhttp)
{
	pllhttp->http_header_finish = 0;
	pllhttp->found_url = 0;
	llhttp_init(&pllhttp->parser, HTTP_REQUEST, &gllhttp_set);
	llhttp_set_priv(&pllhttp->parser, (void*)pllhttp);
}

void child_struct_init(struct child* pchild)
{
	memset(pchild, 0, sizeof(struct child));
}

static void handle_connection_failure(struct conn_s* connptr)
{
	send_http_error_message(connptr);
}

void alloc_buffer(uv_handle_t* handle, size_t suggested_size, uv_buf_t* buf)
{
	buf->base = (char*)malloc(suggested_size);
	buf->len = (ULONG)suggested_size;
}

void write_client_cb(uv_write_t* req, int status)
{
	if (status) {//connection lost
		log_message(LOG_INFO, "Write error %s", uv_strerror(status));
	}
	write_req_t* wr = (write_req_t*)req;
	free(wr->buf.base);
	free(wr);
}

void clean_connection(uv_handle_t* handle)
{
	struct child* pchild = (struct child*)uv_handle_get_data((const uv_handle_t*)handle);
	uv_handle_set_data(handle, NULL);
	free(pchild);
	gChilds_num--;
}

void read_error(uv_stream_t* stream, int nread, const uv_buf_t* buf, int to_server)
{
	if (nread != UV_EOF) {//connection lost
		log_message(LOG_INFO, "Read error %d:%s", to_server, uv_err_name(nread));
		//printf("Read error %d:%s\n", to_server, uv_err_name(nread));
	}
	//one connection includes 2 sockets: one for client/proxy, one for server/proxy.
	//only one struct child (for 2 sockets), release only once
	if(!to_server)
		uv_close((uv_handle_t*)stream, clean_connection);
	free(buf->base);
}

void forward_data(uv_stream_t* stream, unsigned int nread, const uv_buf_t* buf, int to_server)
{
	struct child* pchild = (struct child*)uv_handle_get_data((const uv_handle_t*)stream);
	if (!pchild) //connection lost
		return;
	struct conn_s* connptr = &pchild->conn;

	uv_stream_t* target = NULL;
	if(to_server)
		target = (uv_stream_t*)&connptr->server_socket;
	else
		target = (uv_stream_t*)&connptr->client_socket;

	write_req_t* req = (write_req_t*)malloc(sizeof(write_req_t));
	req->buf = uv_buf_init(buf->base, nread);
	uv_write((uv_write_t*)req, target, &req->buf, 1, write_client_cb);
}

void read_client_data(uv_stream_t* client, ssize_t nread, const uv_buf_t* buf)
{
	if (nread <= 0) {
		read_error(client, (int)nread, buf, 1);
		return;
	}
	forward_data(client, (unsigned int)nread, buf, 1);
}

void read_server_data(uv_stream_t* server, ssize_t nread, const uv_buf_t* buf)
{
	if (nread <= 0) {
		read_error(server, (int)nread, buf, 0);
		return;
	}

	forward_data(server, (unsigned int)nread, buf, 0);
}

/*
 * begin relaying the bytes between the two connections. 
 * We continue to use the buffering code
 * since we want to be able to buffer a certain amount for slower
 * connections (as this was the reason why I originally modified
 * tinyproxy oh so long ago...)
 *	- rjkaes
 */
static void relay_connection(struct conn_s* connptr)
{
	uv_stream_t *client, *server;
	client = (uv_stream_t*)&connptr->client_socket;
	server = (uv_stream_t*)&connptr->server_socket;

	uv_read_start(client, alloc_buffer, read_client_data);
	uv_read_start(server, alloc_buffer, read_server_data);
}

//buf = connptr->cbuf, len = count
void stream_write_client(uv_tcp_t* socket, ssize_t count)
{
	struct child* pchild = (struct child*)uv_handle_get_data((const uv_handle_t*)socket);
	struct conn_s* connptr = &pchild->conn;
	uv_stream_t* client = (uv_stream_t*)&connptr->client_socket;
	uv_buf_t* cbuf = &connptr->cbuf;

	write_req_t* req = (write_req_t*)malloc(sizeof(write_req_t));
	cbuf->len = (ULONG)count;
	req->buf = uv_buf_init(cbuf->base, cbuf->len);
	uv_write((uv_write_t*)req, client, &req->buf, 1, write_client_cb);
}

void on_connect_remote_server(uv_connect_t* req, int status)
{
	struct child* pchild = (struct child*)uv_handle_get_data((const uv_handle_t*)req);
	struct conn_s* connptr = &pchild->conn;
	uv_tcp_t* server = &connptr->server_socket;
	uv_handle_set_data((uv_handle_t*)server, pchild);

	int port = connptr->port;
	char* host = pchild->llhttp.url;

	if (status < 0) {
		log_message(LOG_INFO, "connect failed error %s", uv_err_name(status));
		uv_close((uv_handle_t*)server, NULL);
		log_message(LOG_WARNING,
			"opensock: Could not establish a connection to %s:%d", host, port);

		indicate_http_error(connptr, 500, "Unable to connect",
			"detail",
			PACKAGE_NAME " "
			"was unable to connect to the remote web server.",
			"error", strerror(errno), NULL);
		handle_connection_failure(connptr);
		return;
	}

	const char* response = "HTTP/1.1 200 Connection Established\r\n\r\n";
	connptr->cbuf.base = malloc(RW_BUF_SIZE);
	connptr->cbuf.len = RW_BUF_SIZE;
	strcpy(connptr->cbuf.base, response);
	stream_write_client(&connptr->client_socket, strlen(response));
	log_message(LOG_CONN, "Established connection to host \"%s\"", host);

	relay_connection(connptr);
	log_message(LOG_INFO,
		"Closed connection between local client and remote client,"
		"connection number is %d\n", gChilds_num);
}


/*
 * Read all the headers from the stream
 *
 * int uv_read_start(uv_stream_t *stream, uv_alloc_cb alloc_cb, uv_read_cb read_cb)
Read data from an incoming stream.
The uv_read_cb callback will be made several times
until there is no more data to read or uv_read_stop() is called.

return UV_EINVAL when the stream is closing
 */
void read_client_header(uv_stream_t* client, ssize_t nread, const uv_buf_t* buf)
{
	struct child* pchild = (struct child*)uv_handle_get_data((const uv_handle_t*)client);
	struct conn_s* connptr = &pchild->conn;

	if (nread <= 0) {
		if (nread != UV_EOF) {
			fprintf(stderr, "Read error %s\n", uv_err_name((int)nread));
			log_message(LOG_WARNING, "Could not retrieve all the headers from the client");
			indicate_http_error(connptr, 400, "Bad Request",
				"detail",
				"Could not retrieve all the headers from "
				"the client.", NULL);
			update_stats(STAT_BADCONN);
		}
		goto done;
	} //end if (nread <= 0)

	struct llhttps* pllhttp = &pchild->llhttp;
	llhttp_execute(&pllhttp->parser, buf->base, (size_t)nread);
	if (!pllhttp->http_header_finish)
		goto done;
	uv_read_stop(client);

	struct request_s request;
	if (process_request(pchild, &request)) {
		update_stats(STAT_BADCONN);
		handle_connection_failure(connptr);
		goto done;
	}
#ifdef DEBUG
	//only care debug_url
	if (!pllhttp->found_url || _stricmp(debug_url, pchild->llhttp.url))
		goto done;
#endif
	//parse finished, release read buffer
	free(buf->base);
	opensock(pchild, request.host, request.port);
	return;

done:
	//clean
	uv_read_stop(client);
	if(buf->base)
		free(buf->base);
	gChilds_num--;
	uv_close((uv_handle_t*)client, NULL);
	free(pchild);
	update_stats(STAT_CLOSE);
}


/*
   CONNECT ntp.msn.cn:443 HTTP/1.1\r\n
Host: ntp.msn.cn:443\r\n
Proxy-Connection: keep-alive\r\n
User-Agent: Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/140.0.0.0 Safar...
*/
int handle_connection(struct child* pchild)
{
	struct conn_s* connptr = &pchild->conn;
	uv_tcp_t* client_sock = &pchild->conn.client_socket;
	init_llhttp(&pchild->llhttp); //init http parser

	gChilds_num++;
	uv_handle_set_data((uv_handle_t*)client_sock, pchild);
	uv_read_start((uv_stream_t*)client_sock, alloc_buffer, read_client_header);
	return 0;
}

/*
 * This is the main loop accepting new connections.
 * exit_flag: set by ConsoleCtrlHandler(), when user press ctrl + c
 */
void on_new_connection(uv_stream_t* server, int status)
{
	if (status < 0) {
		fprintf(stderr, "New connection error %s\n", uv_strerror(status));
		return;
	}

	// This is the main loop accepting new connections.
	struct child* pchild = malloc(sizeof(struct child));
	child_struct_init(pchild);
	uv_tcp_t *client_sock = &pchild->conn.client_socket;
	uv_tcp_init(gloop, client_sock);

	//wait connection
	if (uv_accept(server, (uv_stream_t*)client_sock)){
		log_message(LOG_ERR, "Accept returned an error...");
		uv_close((uv_handle_t*)client_sock, NULL);
		free(pchild);
		return;
	}

	handle_connection(pchild);
}

int main(int argc, char** argv)
{
	load_config(argc, argv, &gconfig);
	set_http_proxy(gconfig.proxy_addr, gconfig.proxy_port);
	init_llhttp_set();

	gloop = uv_default_loop();
	SetConsoleCtrlHandler(ConsoleCtrlHandler, TRUE);

	//Start listening on the selected port
	if (start_proxy_server(&gconfig, on_new_connection)) {
		//start server failed
		//reset_http_proxy();
		return 1;
	}
	log_message(LOG_INFO, "Starting main loop. Accepting connections.\n");
	printf("Starting main loop. Accepting connections.\n\n");

	uv_run(gloop, UV_RUN_DEFAULT); //main loop

	//clean action
	exit_proxy_server(&gconfig);
	int ret = uv_loop_close(gloop);
	if (ret) {
		log_message(LOG_INFO, "uv_loop_close error: %s", uv_err_name(ret));
	}

	reset_http_proxy();
	return 0;
}
